# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import xgboost as xgb
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt
import sys
import os
import efinance as ef
from pypinyin import pinyin, Style

def to_pinyin(text):
    s = ''
    for i in pinyin(text, style=Style.NORMAL):
        s += ''.join(i)
    return s

# --- 1. 从命令行获取参数 ---
if len(sys.argv) < 3:
    print("错误：参数不足。")
    print("用法: python3 04_train_model.py <股票代码> <frequency>")
    sys.exit(1)

stock_code = sys.argv[1]
frequency = sys.argv[2]

# --- 2. 获取公司名称并转换为拼音 ---
stock_name_pinyin = stock_code
try:
    quote_df = ef.stock.get_quote_history(stock_code, klt=101, end='20300101', limit=2)
    if not quote_df.empty and '股票名称' in quote_df.columns:
        stock_name_chinese = quote_df['股票名称'].iloc[0]
        stock_name_pinyin = to_pinyin(stock_name_chinese)
except Exception as e:
    print(f"警告：未能获取股票名称。错误: {e}")

# --- 3. 定义输入输出文件名 ---
input_filename = f'output/data/stock_data_{stock_code}_{frequency}_with_indicators.csv'
model_path = f'output/models/xgb_model_{stock_code}_{frequency}.json'
output_image_path = f'output/plots/prediction_vs_actual_{stock_code}_{frequency}.png'

# --- 4. 加载数据 ---
try:
    print(f"开始加载处理好的数据: {input_filename}...")
    df = pd.read_csv(input_filename, index_col='date', parse_dates=True)
    print("数据加载完成。")
except FileNotFoundError:
    print(f"错误: 找不到数据文件 {input_filename}。")
    sys.exit(1)

# --- 5. 创建标签 (Label) ---
df['target'] = df['close'].shift(-1)

# --- 6. 准备特征和标签 ---
df.dropna(inplace=True)
feature_columns = [
    'volume', 'sma5', 'sma20', 'ema12', 'ema26', 'macd', 'macdsignal',
    'macdhist', 'rsi', 'upperband', 'middleband', 'lowerband'
]
X = df[feature_columns]
y = df['target']
print(f"准备了 {X.shape[0]} 条数据，每条数据有 {X.shape[1]} 个特征。")

# --- 7. 划分训练集和测试集 ---
split_percentage = 0.8
split_index = int(len(X) * split_percentage)
X_train, X_test = X[:split_index], X[split_index:]
y_train, y_test = y[:split_index], y[split_index:]
print(f"训练集大小: {len(X_train)} 条, 测试集大小: {len(X_test)} 条")

# --- 8. 训练 XGBoost 模型 ---
print("开始训练 XGBoost 模型...")
model = xgb.XGBRegressor(
    n_estimators=1000, learning_rate=0.05, early_stopping_rounds=5,
    objective='reg:squarederror', random_state=42
)
model.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)
print("模型训练完成。")

# --- 9. 进行预测与评估 ---
predictions = model.predict(X_test)
rmse = np.sqrt(mean_squared_error(y_test, predictions))
print(f"模型的均方根误差 (RMSE) 是: {rmse:.4f}")

# --- 10. 可视化结果 ---
print("开始绘制预测结果图...")
plt.figure(figsize=(15, 7))
plt.plot(y_test.index, y_test.values, label='Actual Price', color='blue')
plt.plot(y_test.index, predictions, label='Predicted Price', color='red', linestyle='--')
plt.title(f'{stock_name_pinyin} ({stock_code}) - {frequency.capitalize()} Price Prediction vs Actual')
plt.xlabel('Date')
plt.ylabel('Price')
plt.legend()
plt.grid(True)
# 确保output目录存在
os.makedirs('output/plots', exist_ok=True)
plt.savefig(output_image_path)
print(f"预测结果图已保存到: {output_image_path}")

# --- 11. 保存训练好的模型 ---
# 确保output/models目录存在
os.makedirs('output/models', exist_ok=True)
model.save_model(model_path)
print(f"模型已成功保存到: {model_path}")